import numpy as np
import torch as pt
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
import sys
import os
from sklearn.preprocessing import StandardScaler

np.random.seed(777)
pt.manual_seed(777)
device = pt.device("cuda:0" if pt.cuda.is_available() else "cpu")
print(device)

VER = 'v2.0'
ALPHA = 0.001
N_EPOCHS = 500

# data
ax = plt.axes(projection='3d')
W1 = 3.5
W2 = 4.8
B = -5.7
M = 11
x = np.linspace(-3, 7, M) + np.random.normal(0.0, 1.0, size=(M,))
y = np.linspace(-5, 5, M) + np.random.normal(0.0, 1.0, size=(M,))
z = W1 * x + W2 * y + B + np.random.normal(0, 2.0, size=(M,))
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
x = x.reshape((-1, 1))
y = y.reshape((-1, 1))
z = y.reshape((-1, 1))
scaler = StandardScaler()
x = scaler.fit_transform(x)
y = scaler.fit_transform(y)
z = scaler.fit_transform(z)
xy = np.c_[x, y]
xyt = pt.tensor(xy).float().to(device)
zt = pt.tensor(z).float().to(device)
xr = x.ravel()
yr = y.ravel()
zr = z.ravel()
ax.scatter3D(xr, yr, zr)

# model
class MyLinModel(pt.nn.Module):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.w = pt.nn.Parameter(pt.randn((2, 1)))
        self.b = pt.nn.Parameter(pt.zeros((1, 1)))

    def forward(self, x):
        y = pt.matmul(x, self.w) + self.b
        return y


model = MyLinModel().to(device)
criterion = pt.nn.MSELoss().to(device)
params = model.parameters()
optim = pt.optim.SGD(params=params, lr=ALPHA)

# train
loss_history = np.zeros(N_EPOCHS)
GROUP = int(np.ceil(N_EPOCHS / 20))
for step in range(N_EPOCHS):
    optim.zero_grad()
    ht = model(xyt)
    cost = criterion(ht, zt)
    cost.backward()
    optim.step()
    costv = cost.cpu().item()
    loss_history[step] = costv
    if step % GROUP == 0:
        print(f'#{step + 1}: cost = {costv}')
if step % GROUP != 0:
    print(f'#{step + 1}: cost = {costv}')
print(f'Weights: {model.w}')
print(f'Bias: {model.b}')

# hypothesis
xlimr = np.array([xr.min(), xr.max()])
xlim = xlimr.reshape((-1, 1))
ylimr = np.array([yr.min(), yr.max()])
ylim = ylimr.reshape((-1, 1))
xylim = np.c_[xlim, ylim]
xylimt = pt.tensor(xylim).float().to(device)
with pt.no_grad():  # ATTENTION
    ht = model(xylimt)
print(ht.shape)
ax.plot3D(xlimr, ylimr, ht.cpu().data.numpy().ravel(), 'red')

plt.show()
sys.exit(0)

from torch.utils.tensorboard import SummaryWriter
filename = os.path.basename(__file__)
logdir = os.path.join('_log', filename, VER)
with SummaryWriter(logdir) as sw:
    sw.add_graph(model, input_to_model=xyt)

# from tensorboard import notebook
# notebook.list()
# notebook.start(f'--logdir {logdir} --port 6039')
